# -*- coding: utf-8 -*-
"""
向量存储管理，支持本地文件存储和Milvus数据库存储
"""

import os
import pickle
import time
import uuid
from typing import List, Dict, Any, Optional, Tuple
import numpy as np
from dataclasses import dataclass, asdict

from .config import embedding_config
from .embedding_client import EmbeddingResult


@dataclass
class VectorMetadata:
    """向量元数据"""
    id: str
    text: str
    vector_dim: int
    model: str
    created_at: float
    tags: Optional[List[str]] = None
    extra_info: Optional[Dict[str, Any]] = None


class VectorStore:
    """向量存储管理类"""
    
    def __init__(self, storage_config: Optional[Dict[str, str]] = None):
        """初始化向量存储"""
        if storage_config is None:
            storage_config = embedding_config.get_storage_config()
        
        self.storage_dir = storage_config["storage_dir"]
        self.vector_file = storage_config["vector_file"]
        self.metadata_file = storage_config["metadata_file"]
        
        # 确保存储目录存在
        os.makedirs(self.storage_dir, exist_ok=True)
        
        # 初始化存储
        self.vectors = np.array([])
        self.metadata = []
        self._load_storage()
    
    def _load_storage(self):
        """加载现有存储"""
        try:
            if os.path.exists(self.vector_file):
                self.vectors = np.load(self.vector_file)
                print(f"加载向量数据，形状: {self.vectors.shape}")
            
            if os.path.exists(self.metadata_file):
                with open(self.metadata_file, 'rb') as f:
                    self.metadata = pickle.load(f)
                print(f"加载元数据，数量: {len(self.metadata)}")
        except Exception as e:
            print(f"加载存储失败: {e}")
            self.vectors = np.array([])
            self.metadata = []
    
    def _save_storage(self):
        """保存存储到文件"""
        try:
            if len(self.vectors) > 0:
                np.save(self.vector_file, self.vectors)
            
            if self.metadata:
                with open(self.metadata_file, 'wb') as f:
                    pickle.dump(self.metadata, f)
            
            print(f"存储保存成功，向量数量: {len(self.vectors)}")
        except Exception as e:
            print(f"保存存储失败: {e}")
    
    def add_vectors(self, embedding_results: List[EmbeddingResult], 
                   tags: Optional[List[str]] = None, 
                   extra_info: Optional[Dict[str, Any]] = None) -> List[str]:
        """添加向量到存储"""
        if not embedding_results:
            return []
        
        new_vectors = []
        new_metadata = []
        vector_ids = []
        
        for result in embedding_results:
            # 生成唯一ID
            vector_id = str(uuid.uuid4())
            vector_ids.append(vector_id)
            
            # 准备向量数据
            vector = np.array(result.vector, dtype=np.float32)
            new_vectors.append(vector)
            
            # 准备元数据
            metadata = VectorMetadata(
                id=vector_id,
                text=result.text,
                vector_dim=len(result.vector),
                model=result.model,
                created_at=time.time(),
                tags=tags,
                extra_info=extra_info
            )
            new_metadata.append(metadata)
        
        # 更新存储
        if len(self.vectors) == 0:
            self.vectors = np.array(new_vectors)
        else:
            self.vectors = np.vstack([self.vectors, new_vectors])
        
        self.metadata.extend(new_metadata)
        
        # 保存到文件
        self._save_storage()
        
        print(f"成功添加 {len(embedding_results)} 个向量")
        return vector_ids
    
    def get_vector(self, vector_id: str) -> Optional[Tuple[np.ndarray, VectorMetadata]]:
        """根据ID获取向量和元数据"""
        for i, metadata in enumerate(self.metadata):
            if metadata.id == vector_id:
                return self.vectors[i], metadata
        return None
    
    def search_similar(self, query_vector: List[float], top_k: int = 5, 
                      threshold: float = 0.0) -> List[Tuple[str, float, VectorMetadata]]:
        """搜索相似向量"""
        if len(self.vectors) == 0:
            return []
        
        query_array = np.array(query_vector, dtype=np.float32)
        
        # 计算余弦相似度
        similarities = self._cosine_similarity(query_array, self.vectors)
        
        # 获取top-k结果
        top_indices = np.argsort(similarities)[::-1][:top_k]
        
        results = []
        for idx in top_indices:
            similarity = similarities[idx]
            if similarity >= threshold:
                metadata = self.metadata[idx]
                results.append((metadata.id, similarity, metadata))
        
        return results
    
    def search_by_text(self, query_text: str, top_k: int = 5, 
                      threshold: float = 0.0) -> List[Tuple[str, float, VectorMetadata]]:
        """根据文本搜索相似向量"""
        # 这里需要先对查询文本进行向量化
        # 为了简化，这里假设已经有向量化的客户端
        from .embedding_client import EmbeddingClient
        
        client = EmbeddingClient()
        try:
            query_result = client.embed_text(query_text)
            return self.search_similar(query_result.vector, top_k, threshold)
        except Exception as e:
            print(f"文本向量化失败: {e}")
            return []
    
    def delete_vector(self, vector_id: str) -> bool:
        """删除向量"""
        for i, metadata in enumerate(self.metadata):
            if metadata.id == vector_id:
                # 删除向量和元数据
                self.vectors = np.delete(self.vectors, i, axis=0)
                del self.metadata[i]
                
                # 保存更新
                self._save_storage()
                print(f"成功删除向量: {vector_id}")
                return True
        
        print(f"未找到向量: {vector_id}")
        return False
    
    def get_all_vectors(self) -> List[Tuple[str, np.ndarray, VectorMetadata]]:
        """获取所有向量"""
        return [(metadata.id, self.vectors[i], metadata) 
                for i, metadata in enumerate(self.metadata)]
    
    def get_stats(self) -> Dict[str, Any]:
        """获取存储统计信息"""
        return {
            "total_vectors": len(self.vectors),
            "vector_dimension": self.vectors.shape[1] if len(self.vectors) > 0 else 0,
            "storage_size_mb": os.path.getsize(self.vector_file) / (1024 * 1024) if os.path.exists(self.vector_file) else 0,
            "models": list(set(m.model for m in self.metadata)),
            "created_date_range": {
                "earliest": min(m.created_at for m in self.metadata) if self.metadata else None,
                "latest": max(m.created_at for m in self.metadata) if self.metadata else None
            }
        }
    
    def _cosine_similarity(self, query: np.ndarray, vectors: np.ndarray) -> np.ndarray:
        """计算余弦相似度"""
        # 归一化向量
        query_norm = query / np.linalg.norm(query)
        vectors_norm = vectors / np.linalg.norm(vectors, axis=1, keepdims=True)
        
        # 计算余弦相似度
        similarities = np.dot(vectors_norm, query_norm)
        return similarities


class MilvusVectorStore:
    """Milvus向量数据库存储类（预留接口）"""
    
    def __init__(self, connection_params: Dict[str, Any]):
        """初始化Milvus连接"""
        self.connection_params = connection_params
        self.collection_name = "text_vectors"
        # 这里可以添加Milvus客户端初始化代码
        print("Milvus存储功能待实现")
    
    def add_vectors(self, embedding_results: List[EmbeddingResult]) -> List[str]:
        """添加向量到Milvus"""
        # 预留Milvus实现
        print("Milvus添加向量功能待实现")
        return []
    
    def search_similar(self, query_vector: List[float], top_k: int = 5) -> List[Tuple[str, float]]:
        """在Milvus中搜索相似向量"""
        # 预留Milvus实现
        print("Milvus搜索功能待实现")
        return []
